from Dataset.load_by_time import DataSet
from datetime import datetime
from Sketches.vertex_bias_reservoir_sampling import ReservoirSampling
from Sketches.min_hash import MinHash
from Dataset.util import Vertex
from Estimation.common_neighbor import CommonNeighbor
from Estimation.adamic_adar import AdamicAdar
from Estimation.jaccard import Jaccard


def estimation():
    time1 = datetime.now()
    print(time1)
    # --- parameter ---
    Sketch = None
    para = None
    Measure = None
    name = 'wiki'
    end = "2002-01-01 06:22:33"
    measure_type = 'aa'

    # --- init ---
    if measure_type == 'cn':
        Sketch = ReservoirSampling
        Measure = CommonNeighbor
        para = {"reverse": False, "budget": 100}
    elif measure_type == 'aa':
        Sketch = ReservoirSampling
        Measure = AdamicAdar
        para = {"reverse": True, "budget": 100}
    elif measure_type == 'jaccard':
        Sketch = MinHash
        Measure = Jaccard
        para = {"perm_num": 20}

    dataset = DataSet(end, name=name)
    dataset.next_time = dataset.start_time
    vertices = {}
    core = []
    train_edges = []
    test_edges = []
    predictor_edge = []
    measures = {}
    predicted = 0
    edge_count = 0

    while True:
        if dataset.next_time < dataset.end_time:
            trunk = dataset.get_next()
            if dataset.next_time < dataset.division_time:
                for edge in trunk:
                    edge_count += 1
                    vertex1 = vertices.get(edge.vertex[0])
                    vertex2 = vertices.get(edge.vertex[1])
                    if vertex1 is None:
                        vertex1 = Vertex(edge.vertex[0])
                        vertex1.sketch = Sketch(vertex1.vertex, para)
                        vertices[vertex1.vertex] = vertex1
                    vertex1.add_degree()
                    if vertex2 is None:
                        vertex2 = Vertex(edge.vertex[1])
                        vertex2.sketch = Sketch(vertex2.vertex, para)
                        vertices[vertex2.vertex] = vertex2
                    vertex2.add_degree()
                    vertex1.sketch.add_vertex(vertex2)
                    vertex2.sketch.add_vertex(vertex1)
                    edge_id = edge.get_identify()
                    if edge_id not in train_edges:
                        train_edges.append(edge_id)
                    '''
                    vertex_pair = []
                    for vid in edge.vertex:
                        if vid in vertexes:
                            vertex = vertexes[vid]
                        else:
                            vertex = Vertex(vid)
                            vertex.sketch = Sketch(vid)
                            vertexes[vid] = vertex
                        vertex.add_degree()
                    '''
                    '''
                        vertex_pair.append(vertex)
                    for vertex in vertex_pair:
                        changed, vertex_u = vertex.sketch.add_vertex(vertex_pair)
                        if changed:
                            for vertex_v_id in vertex_u.inverse_sampling:
                                if vertex_v_id != vertex.vertex:
                                    contained = vertex.add_cn(vertexes[vertex_v_id])
                                    cn_id = str(min(vertex.vertex, vertex_v_id)) + '+' + str(max(vertex.vertex,
                                                                                                 vertex_v_id))
                                    if contained:
                                        cn = measures[cn_id]
                                        cn.calc_value(vertex, vertexes[vertex_v_id])
                                    else:
                                        cn = Measure(vertex, vertexes[vertex_v_id])
                                        measures[cn_id] = cn
                    '''
            else:
                for edge in trunk:
                    edge_count += 1
                    for vertex in edge.vertex:
                        if vertex in vertices:
                            if not vertices[vertex] in core:
                                core.append(vertices[vertex])
                    '''
                    edge_id = edge.get_identify()
                    if edge_id not in test_edges:
                        test_edges.append(edge_id)
                    '''
                    test_edges.append(set(edge.vertex))

            print(dataset.current_time)
        else:
            break

    core_count = len(core)
    train_count = len(train_edges)

    for start_index in range(core_count):
        for end_index in range(start_index, core_count):
            if not ({core[start_index], core[end_index]} in train_edges) | (core[start_index] == core[end_index]):
                m_id = str(min(core[start_index], core[end_index])) + '+' + str(max(core[start_index], core[end_index]))
                if m_id not in measures:
                    measure = Measure(core[start_index], core[end_index], vertices)
                    if measure.value != 0:
                        measures[m_id] = measure

    measures = list(measures.values())
    measures.sort(reverse=True)

    for index in range(len(test_edges)):
        if index < len(measures):
            predictor_edge.append(set(measures[index].vertex))

    for edge in predictor_edge:
        if edge in test_edges:
            predicted += 1
    abs_accuracy = predicted / len(predictor_edge)
    random_accuracy = len(test_edges) / (core_count * (core_count - 1) / 2 - train_count)
    relative_accuracy = abs_accuracy / random_accuracy
    print(abs_accuracy)
    print(relative_accuracy)
    print(predicted)

    time2 = datetime.now()
    print(time2)
    time12 = (time2 - time1)
    print(time12)


estimation()
